#!/usr/bin/python

import urllib                   # to download files
import sys                      # to write messages to buffers
import getopt                   # to handle command line options
import subprocess               # to use external unzippers
import os.path                  # to manipulate filenames
import shutil                   # file copying

from tcclimFunctions import *   # get common functions 

## FULL DATA SCRIPT HAMPERED BY THE LACK OF LZW FILE DECOMPRESSION FOR PYTHON
## BUT 7ZIP provides a windows solution for .Z and pax archives - maybe also Mac/UNIX, 
## but pax/uncompress also available for those platforms

## ~ 15 minutes including download time on a standard Windows 7 desktop

def usage():
  
  use = """
  
  tcclimDownloadTYNSC2.py:
  
  Downloads component data files for the TYN SC 2.03 climate
  datasets and creates transposed data files for use by the
  tcclimCruTS2.py program. The downloaded files will be saved
  in the folder 'tcclim_tyn_sc_2_03'.
  
  Basic use:
     tcclimDownloadTYNSC2.py
  
  Options:
     -h or --help: show this help information
     -k or --keep: retain the original untransposed data files
     -7 or --7zip: path to the 7zip program on Windows
     -u or --uncompress: path to the uncompress program on Mac or Linux
     -p or --pax: path to the pax program on Mac or Linux
  
  This program requires the use of external programs to extract
  the original data from compressed file formats. On Windows, 
  the open-source program 7zip is used and must be installed 
  before using this program. By default, the 7zip program is 
  expected to be found at 'C:\Program Files\7Zip\7Z.exe', but 
  this may be changed using the --7zip option. On Mac or Linux,
  pax and uncompress are used and the paths to these may also
  be specified.
  
  """
  
  sys.stdout.write(use)


def main():
    
    # sort out the options and argument
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hks:u:p:", ["help","keep","sevenzip","uncompress","pax"])
    except getopt.GetoptError, err:
        # print help information and exit:
        sys.stderr.write("\n  " + str(err) + '\n') # will print something like "option -a not recognized"
        usage()
        sys.exit()
    
    # initial setups
    keep = False
    sevenzip = r"C:\Program Files\7-Zip\7Z.exe"
    uncompress = 'uncompress'
    pax = "pax"
    
    for o, a in opts:
        if o in ('-h', '--help'):
            usage()
            sys.exit()
        elif o in ('k', '--keep'):
            keep = True
        elif o in ('s', '--sevenzip'):
            sevenzip = a
        elif o in ('u', '--uncompress'):
            uncompress = a
        elif o in ('', '--pax'):
            pax = a
        else:
            assert False, "Unhandled option"
            usage()
            sys.exit()
    
    # preflight checks
    # - test for utilities
    if os.name == 'posix':
        # use 'which' to test for uncompress
        try:  
            proc = subprocess.Popen('which ' + uncompress ,
                                    shell=True, stdout=subprocess.PIPE,
                                    stderr=subprocess.PIPE)
            procComm = proc.communicate() # get the stdout and stderr
            # the stdout of procComm will be an empty string if the executable is not found
            if procComm[0] == '':
                sys.stderr.write('\nuncompress not found: ' + uncompress + '\n\n')
                usage()
                sys.exit() 
        except OSError, e:
            sys.stderr.write("Execution failed:" + e)
        
        # use 'which' to test for pax
        try:  
            proc = subprocess.Popen('which ' + pax ,
                                    shell=True, stdout=subprocess.PIPE,
                                    stderr=subprocess.PIPE)
            procComm = proc.communicate() # get the stdout and stderr
            # the stdout of procComm will be an empty string if the executable is not found
            if procComm[0] == '':
                sys.stderr.write('\npax not found: ' + pax + '\n\n')
                usage()
                sys.exit() 
        except OSError, e:
            sys.stderr.write("Execution failed:" + e)
        
    elif os.name == 'nt':
        if not os.path.exists(sevenzip):
            sys.stderr.write('\n7Zip program not found\n\n')
            usage()
            sys.exit()
    
    ## create data directories
    if not os.path.exists('tyn_sc_2_03'):
        os.mkdir('tyn_sc_2_03')
    else:
        sys.stderr.write('Download folder \'tyn_sc_2_03\' already exists.\n\n')
        sys.exit()
    
    if not os.path.exists('tcclim_tyn_sc_2_03'):
        os.makedirs('tcclim_tyn_sc_2_03')
    else:
        sys.stderr.write('\nTransposed file directory \'tcclim_tyn_sc_2_03\' already exists...\n\n')
        sys.exit()
    
    ##########################
    ## DOWNLOAD TYNSC  2.10 ##
    ##########################
    
    tynsc2repos = 'http://www.cru.uea.ac.uk/cru/data/hrg/timm/tyn_sc_2.03/data/'
    
    # files to download
    tynsc2files = ['archive.hd', 'archive.ann', 'obs.clim6190.globe.lan.cld.Z',
                   'obs.clim6190.globe.lan.dtr.Z', 'obs.clim6190.globe.lan.pre.Z',
                   'obs.clim6190.globe.lan.tmp.Z', 'obs.clim6190.globe.lan.vap.Z',
                   'iavar.1901-2000.cld.Z', 'iavar.1901-2000.dtr.Z', 'iavar.1901-2000.pre.Z',
                   'iavar.1901-2000.tmp.Z', 'iavar.1901-2000.vap.Z', 'scenarios.txt']
    
    sys.stdout.write('\n\nDownloading TYN SC 2.03      [.....................]\n')
    
    for eachF in tynsc2files:
        remotename = tynsc2repos + eachF
        localname  = os.path.join('tyn_sc_2_03', eachF)
        sys.stdout.write(eachF.ljust(29) + '[')
        sys.stdout.flush()
        try:
            urllib.urlretrieve(remotename, localname, progressTicker)
        except IOError, e[1][1]:
            sys.stderr.write("\n\nFile " + eachF + " could not be downloaded:\n" + e + "\nCheck network connection.\n\n")
            sys.exit()
        sys.stdout.write(']\n')
    
    sys.stdout.write('\nTYN SC 2.03 dataset download complete.\n\n')
    
    ##########################
    ## EXTRACT TYN SC 2.03  ##
    ##########################
    
    os.chdir('tyn_sc_2_03')
    
    # slight trick - one of the archive files (pax) contains
    # compressed (.Z) files - so two steps:
    #    1) Extract the pax archives 
    #    2) Uncompress the downloded and extracted .Z files
    
    compressFiles =  ['hd.cgcm2.a1fi.2080s.cld.Z', 'hd.cgcm2.a1fi.2080s.dtr.Z',
                      'hd.cgcm2.a1fi.2080s.pre.Z', 'hd.cgcm2.a1fi.2080s.tmp.Z',
                      'hd.cgcm2.a1fi.2080s.vap.Z', 'hd.cgcm2.a2a.2080s.cld.Z',
                      'hd.cgcm2.a2a.2080s.dtr.Z', 'hd.cgcm2.a2a.2080s.pre.Z',
                      'hd.cgcm2.a2a.2080s.tmp.Z', 'hd.cgcm2.a2a.2080s.vap.Z',
                      'hd.cgcm2.b1.2080s.cld.Z', 'hd.cgcm2.b1.2080s.dtr.Z',
                      'hd.cgcm2.b1.2080s.pre.Z', 'hd.cgcm2.b1.2080s.tmp.Z',
                      'hd.cgcm2.b1.2080s.vap.Z', 'hd.cgcm2.b2a.2080s.cld.Z',
                      'hd.cgcm2.b2a.2080s.dtr.Z', 'hd.cgcm2.b2a.2080s.pre.Z',
                      'hd.cgcm2.b2a.2080s.tmp.Z', 'hd.cgcm2.b2a.2080s.vap.Z',
                      'hd.csiro2.a1a.2080s.cld.Z', 'hd.csiro2.a1a.2080s.dtr.Z',
                      'hd.csiro2.a1a.2080s.pre.Z', 'hd.csiro2.a1a.2080s.tmp.Z',
                      'hd.csiro2.a1a.2080s.vap.Z', 'hd.csiro2.a2a.2080s.cld.Z',
                      'hd.csiro2.a2a.2080s.dtr.Z', 'hd.csiro2.a2a.2080s.pre.Z',
                      'hd.csiro2.a2a.2080s.tmp.Z', 'hd.csiro2.a2a.2080s.vap.Z',
                      'hd.csiro2.b1a.2080s.cld.Z', 'hd.csiro2.b1a.2080s.dtr.Z',
                      'hd.csiro2.b1a.2080s.pre.Z', 'hd.csiro2.b1a.2080s.tmp.Z',
                      'hd.csiro2.b1a.2080s.vap.Z', 'hd.csiro2.b2a.2080s.cld.Z',
                      'hd.csiro2.b2a.2080s.dtr.Z', 'hd.csiro2.b2a.2080s.pre.Z',
                      'hd.csiro2.b2a.2080s.tmp.Z', 'hd.csiro2.b2a.2080s.vap.Z',
                      'hd.echam4.a2.lan.cld.Z', 'hd.echam4.a2.lan.dtr.Z', 'hd.echam4.a2.lan.pre.Z',
                      'hd.echam4.a2.lan.tmp.Z', 'hd.echam4.a2.lan.vap.Z', 'hd.echam4.b2.lan.cld.Z',
                      'hd.echam4.b2.lan.dtr.Z', 'hd.echam4.b2.lan.pre.Z', 'hd.echam4.b2.lan.tmp.Z',
                      'hd.echam4.b2.lan.vap.Z', 'hd.had3.a1fi.2080s.cld.Z',
                      'hd.had3.a1fi.2080s.dtr.Z', 'hd.had3.a1fi.2080s.pre.Z',
                      'hd.had3.a1fi.2080s.tmp.Z', 'hd.had3.a1fi.2080s.vap.Z',
                      'hd.had3.a2m.2080s.cld.Z', 'hd.had3.a2m.2080s.dtr.Z',
                      'hd.had3.a2m.2080s.pre.Z', 'hd.had3.a2m.2080s.tmp.Z',
                      'hd.had3.a2m.2080s.vap.Z', 'hd.had3.b1.2080s.cld.Z', 'hd.had3.b1.2080s.dtr.Z',
                      'hd.had3.b1.2080s.pre.Z', 'hd.had3.b1.2080s.tmp.Z', 'hd.had3.b1.2080s.vap.Z',
                      'hd.had3.b2m.2080s.cld.Z', 'hd.had3.b2m.2080s.dtr.Z',
                      'hd.had3.b2m.2080s.pre.Z', 'hd.had3.b2m.2080s.tmp.Z',
                      'hd.had3.b2m.2080s.vap.Z', 'hd.pcm.a1fi.2080s.cld.Z',
                      'hd.pcm.a1fi.2080s.dtr.Z', 'hd.pcm.a1fi.2080s.pre.Z',
                      'hd.pcm.a1fi.2080s.tmp.Z', 'hd.pcm.a1fi.2080s.vap.Z',
                      'hd.pcm.a2a.2080s.cld.Z', 'hd.pcm.a2a.2080s.dtr.Z', 'hd.pcm.a2a.2080s.pre.Z',
                      'hd.pcm.a2a.2080s.tmp.Z', 'hd.pcm.a2a.2080s.vap.Z', 'hd.pcm.b1.2080s.cld.Z',
                      'hd.pcm.b1.2080s.dtr.Z', 'hd.pcm.b1.2080s.pre.Z', 'hd.pcm.b1.2080s.tmp.Z',
                      'hd.pcm.b1.2080s.vap.Z', 'hd.pcm.b2a.2080s.cld.Z', 'hd.pcm.b2a.2080s.dtr.Z',
                      'hd.pcm.b2a.2080s.pre.Z', 'hd.pcm.b2a.2080s.tmp.Z', 'hd.pcm.b2a.2080s.vap.Z',
                      'iavar.1901-2000.cld.Z', 'iavar.1901-2000.dtr.Z', 'iavar.1901-2000.pre.Z',
                      'iavar.1901-2000.tmp.Z', 'iavar.1901-2000.vap.Z', 
                      'obs.clim6190.globe.lan.cld.Z', 'obs.clim6190.globe.lan.dtr.Z',
                      'obs.clim6190.globe.lan.pre.Z', 'obs.clim6190.globe.lan.tmp.Z',
                      'obs.clim6190.globe.lan.vap.Z']
    
    # External uncompression differs by platform
    # use os.name to get 'nt' (Windows) versus 'posix' (Darwin/Linux)
    # - could use sys.platform but more variation in mac/linux returns
    if os.name == 'posix':
        # use pax for the two archive files
        for eachF in tynsc2files[0:2]:
            sys.stdout.write('Uncompressing ' + eachF + '\n')
            sys.stdout.flush()
            try:  
               proc = subprocess.Popen(pax + ' -r -f ' + eachF,
                                       shell=True, stderr=subprocess.PIPE,
                                       stdout=subprocess.PIPE)
               procWait = proc.wait()        # hang on until the extraction ends
               procComm = proc.communicate() # get the stdout and stderr
            except OSError, e:
               sys.stderr.write("Execution failed:" + e)
        
        # use compress for the other Z files
        for eachF in compressFiles:
            sys.stdout.write('Uncompressing ' + eachF + '\n')
            sys.stdout.flush()
            try:  
               proc = subprocess.Popen(uncompress + ' ' + eachF,
                                       shell=True, stderr=subprocess.PIPE,
                                       stdout=subprocess.PIPE)
               procWait = proc.wait()        # hang on until the extraction ends
               procComm = proc.communicate() # get the stdout and stderr
            except OSError, e:
               sys.stderr.write("Execution failed:" + e)
        
    elif os.name == 'nt':
        # use 7zip
        for eachF in tynsc2files[0:2]:
            sys.stdout.write('Uncompressing ' + eachF + '\n')
            sys.stdout.flush()
            try:  
               proc = subprocess.Popen('"' + sevenzip + '" e ' + eachF,
                                       shell=True, stderr=subprocess.PIPE,
                                       stdout=subprocess.PIPE)
               procWait = proc.wait()        # hang on until the extraction ends
               procComm = proc.communicate() # get the stdout and stderr
            except OSError, e:
               sys.stderr.write("Execution failed:" + e)
        
        for eachF in compressFiles:
            sys.stdout.write('Uncompressing ' + eachF + '\n')
            sys.stdout.flush()
            try:  
               proc = subprocess.Popen('"' + sevenzip + '" e ' + eachF,
                                       shell=True, stderr=subprocess.PIPE,
                                       stdout=subprocess.PIPE)
               procWait = proc.wait()        # hang on until the extraction ends
               procComm = proc.communicate() # get the stdout and stderr
            except OSError, e:
               sys.stderr.write("Execution failed:" + e)
    
    ###########################
    ## TRANSPOSE TYN SC 2.03 ##
    ###########################
    
    ## TRANSPOSE AN EXISTING COPY OF THE TYN SC 2.03 DATASET
    ## - check directories
    
    os.chdir('../')
    sys.stdout.write('\nTransposing data to \'tcclim_tyn_sc_2_03\'.\n\n')
    
    ## naming of files
    
    modelsHD = ['cgcm2', 'csiro2', 'echam4', 'had3', 'pcm']
    
    scenariosHD = {'cgcm2':  ['cgcm2.a1fi.2080s','cgcm2.a2a.2080s','cgcm2.b1.2080s','cgcm2.b2a.2080s'],
                   'csiro2': ['csiro2.a1a.2080s','csiro2.a2a.2080s','csiro2.b1a.2080s','csiro2.b2a.2080s'],
                   'echam4': ['','echam4.a2.lan','echam4.b2.lan',''],
                   'had3':   ['had3.a1fi.2080s','had3.a2m.2080s','had3.b1.2080s','had3.b2m.2080s'],
                   'pcm':    ['pcm.a1fi.2080s','pcm.a2a.2080s','pcm.b1.2080s','pcm.b2a.2080s']}
    
    variables = ['dtr','pre','tmp','vap','cld']
    
    ################################################
    ## translate hd: projected change (2070-2099) ##
    ## (4.5 models * 4 scenarios * 5 variables)   ##
    ################################################
    
    for mod in modelsHD:
        for scen in scenariosHD[mod]:
            # skip the empty echam scenarios
            if scen == '': continue  
            for var in variables:
                
                # base filename
                fname = 'hd.' + scen + '.' + var
                
                # transpose
                infile  = 'tyn_sc_2_03/' + fname
                outfile = 'tcclim_tyn_sc_2_03/tr.' + fname
                sys.stdout.write('Transposing ' + fname + '\n')
                try:
                    transposeGrim(infile, outfile)
                except:
                    sys.stderr.write( 'Skipping missing file: ' + fname + '\n')
    
    ################################################
    ## translate obs: mean observed 1961-1990     ##
    ## (5 variables)                              ##
    ################################################
    
    for var in variables:
        
        # base filename
        fname = 'obs.clim6190.globe.lan.' + var
        
        # transpose
        infile  = 'tyn_sc_2_03/' + fname
        outfile = 'tcclim_tyn_sc_2_03/tr.' + fname
        sys.stdout.write( 'Transposing ' + fname + '\n')
        try:
            transposeGrim(infile, outfile)
        except:
            sys.stderr.write( 'Skipping missing file: ' + fname + '\n')
    
    ################################################
    ## translate iavar: interannual+decadal var   ##
    ## (5 variables)                              ##
    ################################################
    
    for var in variables:
        
        # base filename
        fname = 'iavar.1901-2000.' + var
        
        # transpose
        infile  = 'tyn_sc_2_03/' + fname
        outfile = 'tcclim_tyn_sc_2_03/tr.' + fname
        sys.stdout.write( 'Transposing ' + fname + '\n')
        try:
            transposeGrim(infile, outfile)
        except:
            sys.stderr.write( 'Skipping missing file: ' + fname + '\n')
    
    ###################################################
    ## create a file containing grid-refs in the     ##
    ## order in which they appear in the grim files  ##
    ###################################################
    
    # arbitrary choice of small (not iavar) grim file
    getGridRefIndexGrim('tyn_sc_2_03/obs.clim6190.globe.lan.pre', 'tcclim_tyn_sc_2_03/transposedGridRefIndex.txt')
    
    ###################################
    ## copy ann files into transpose ##
    ## without any change            ##
    ###################################
    
    modelsAnn = ['cgcm2', 'csiro2', 'echam4', 'hadcm3', 'pcm']
    scenariosAnn = ['a1fi','a2','b1','b2']
    
    sys.stdout.write( 'Copying .ann files.\n')
    
    for mod in modelsAnn:
        for scen in scenariosAnn:
            ann = 'tyn_sc_2_03/' + mod + '-' + scen + '.ann'
            try:
                shutil.copy2(ann, 'tcclim_tyn_sc_2_03')
            except:
                sys.stderr.write( 'Skipping missing file: ' + ann + '\n')
    
    
    
    # scan last input grim file for grid refs
    getGridRefIndexGrim(infile, 'tcclim_tyn_sc_2_03/transposedGridRefIndex.txt')
    
    if keep:
        sys.stdout.write( '\n\nCRU TS 2.10 transposition completed.\nOriginal data directory retained.\n\n')
    else:
        shutil.rmtree('tyn_sc_2_03/')
        sys.stdout.write( '\n\nCRU TS 2.10 transposition completed.\nOriginal data directory deleted.\n\n')


if __name__ == "__main__":
    main()

